from sklearn.model_selection import train_test_split
import cv2 as cv
from util import get,dump1,preprocess_image
import os  
import numpy as np
import glob
from tqdm import tqdm
import time

# 1. 读取配置文件中的信息
train_dir = get("train") # 获取 训练数据路径
char_styles = get("char_styles") # 获取 字符样式列表，注意: 必须是列标
new_size = get("new_size") # 获取 新图像大小元组, 注意: 必须包含h和w

# 2. 生成X,y 
print("# 读取训练数据并进行预处理，") 
X=[]
y=[]
l1=['train_篆*','train_隶*','train_草*','train_行*','train_楷*']
for i in l1:
    img_files=glob.glob(f'{train_dir}/{i}')
    for t in tqdm(np.linspace(0,len(img_files)-1,1000,dtype=int),desc=f"处理{(i.split('_')[1]).split('*')[0]}书的进度",unit='it'):
        x=preprocess_image(img_files[t],new_size)
        X.append(x)
        y.append(l1.index(i))
        time.sleep(0.001)
        #进度条



X=np.array(X,dtype=np.float64)
y=np.array(y,dtype=np.int64)

# 3. 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.2)

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
dump1((X_train,X_test,y_train,y_test),"X_train, X_test, y_train, y_test", f'{get("Xy_root")}/Xy')
